import tkinter as tk
from tkinter import ttk, scrolledtext
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
import brain_model
from matplotlib.colors import LinearSegmentedColormap
import time

class BrainVisualizer:
    """人工脑模型的图形化界面"""
    def __init__(self, root):
        self.root = root
        self.root.title("人工脑模型可视化工具")
        self.root.geometry("1200x800")
        
        # 创建脑模型
        self.brain = brain_model.ArtificialBrainModel()
        
        # 存储激活历史
        self.activation_history = {name: [] for name in self.brain.regions}
        self.time_steps = []
        self.current_step = 0
        
        # 创建主界面
        self._create_ui()
        
        # 初始化演示数据
        self._init_demo_data()
        
    def _create_ui(self):
        """创建用户界面"""
        # 创建主分隔窗口
        main_paned = ttk.PanedWindow(self.root, orient=tk.HORIZONTAL)
        main_paned.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)
        
        # 左侧控制面板
        control_frame = ttk.LabelFrame(main_paned, text="控制面板")
        main_paned.add(control_frame, weight=1)
        
        # 右侧可视化区域
        vis_frame = ttk.LabelFrame(main_paned, text="可视化区域")
        main_paned.add(vis_frame, weight=4)
        
        # 设置控制面板
        self._setup_control_panel(control_frame)
        
        # 设置可视化区域
        self._setup_visualization_panel(vis_frame)
        
    def _setup_control_panel(self, parent):
        """设置控制面板"""
        # 刺激输入区域
        stimulus_frame = ttk.LabelFrame(parent, text="刺激输入")
        stimulus_frame.pack(fill=tk.X, padx=5, pady=5)
        
        ttk.Label(stimulus_frame, text="刺激类型:").pack(anchor=tk.W, padx=5, pady=2)
        self.stimulus_type = tk.StringVar(value="visual")
        ttk.Radiobutton(stimulus_frame, text="视觉", variable=self.stimulus_type, value="visual").pack(anchor=tk.W, padx=20)
        ttk.Radiobutton(stimulus_frame, text="听觉", variable=self.stimulus_type, value="auditory").pack(anchor=tk.W, padx=20)
        ttk.Radiobutton(stimulus_frame, text="触觉", variable=self.stimulus_type, value="touch").pack(anchor=tk.W, padx=20)
        
        ttk.Label(stimulus_frame, text="刺激内容:").pack(anchor=tk.W, padx=5, pady=2)
        self.stimulus_content = tk.StringVar(value="")
        ttk.Entry(stimulus_frame, textvariable=self.stimulus_content).pack(fill=tk.X, padx=5, pady=5)
        
        ttk.Button(stimulus_frame, text="发送刺激", command=self._send_stimulus).pack(pady=5)
        
        # 演示区域
        demo_frame = ttk.LabelFrame(parent, text="演示示例")
        demo_frame.pack(fill=tk.X, padx=5, pady=5)
        
        ttk.Button(demo_frame, text="刀子学习演示", command=self._run_knife_demo).pack(fill=tk.X, padx=5, pady=5)
        ttk.Button(demo_frame, text="语言学习演示", command=self._run_language_demo).pack(fill=tk.X, padx=5, pady=5)
        
        # 系统状态
        status_frame = ttk.LabelFrame(parent, text="系统状态")
        status_frame.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)
        
        self.status_text = scrolledtext.ScrolledText(status_frame, wrap=tk.WORD)
        self.status_text.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)
        self.status_text.config(state=tk.DISABLED)
        
        # 底部控制按钮
        control_buttons = ttk.Frame(parent)
        control_buttons.pack(fill=tk.X, padx=5, pady=5)
        
        ttk.Button(control_buttons, text="重置系统", command=self._reset_system).pack(side=tk.LEFT, padx=5)
        ttk.Button(control_buttons, text="清空记录", command=self._clear_log).pack(side=tk.LEFT, padx=5)
        
    def _setup_visualization_panel(self, parent):
        """设置可视化区域"""
        # 创建内部分隔窗口
        vis_paned = ttk.PanedWindow(parent, orient=tk.VERTICAL)
        vis_paned.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)
        
        # 上方为区域激活热力图
        heatmap_frame = ttk.LabelFrame(vis_paned, text="神经区域激活热力图")
        vis_paned.add(heatmap_frame, weight=3)
        
        # 热力图
        self.fig_heatmap, self.ax_heatmap = plt.subplots(figsize=(8, 4))
        self.canvas_heatmap = FigureCanvasTkAgg(self.fig_heatmap, master=heatmap_frame)
        self.canvas_heatmap.get_tk_widget().pack(fill=tk.BOTH, expand=True)
        
        # 下方为网络连接图
        connections_frame = ttk.LabelFrame(vis_paned, text="区域间连接")
        vis_paned.add(connections_frame, weight=3)
        
        # 创建区域图
        self.fig_connections, self.ax_connections = plt.subplots(figsize=(8, 4))
        self.canvas_connections = FigureCanvasTkAgg(self.fig_connections, master=connections_frame)
        self.canvas_connections.get_tk_widget().pack(fill=tk.BOTH, expand=True)
        
        # 底部为时间序列图
        timeseries_frame = ttk.LabelFrame(vis_paned, text="神经区域活动时间序列")
        vis_paned.add(timeseries_frame, weight=2)
        
        # 时间序列图
        self.fig_timeseries, self.ax_timeseries = plt.subplots(figsize=(8, 3))
        self.canvas_timeseries = FigureCanvasTkAgg(self.fig_timeseries, master=timeseries_frame)
        self.canvas_timeseries.get_tk_widget().pack(fill=tk.BOTH, expand=True)
        
        # 初始化绘图
        self._init_plots()
        
    def _init_plots(self):
        """初始化绘图"""
        # 初始化热力图
        self._update_heatmap()
        
        # 初始化连接图
        self._update_connections()
        
        # 初始化时间序列图
        self._update_timeseries()
        
    def _init_demo_data(self):
        """初始化演示数据"""
        # 创建刀子模式
        self.knife_visual = np.random.rand(200) * 0.5 + 0.5
        self.knife_word = np.zeros(300)
        for i, char in enumerate("刀子"):
            self.knife_word[i] = ord(char) / 255
            
        self.pain_pattern = np.random.rand(80) * 0.8 + 0.2
        self.withdraw_pattern = np.random.rand(120) * 0.9
        
        # 创建"我们"概念相关模式
        self.wo_word = np.zeros(300)  # "我"的声音
        for i, char in enumerate("我"):
            self.wo_word[i] = ord(char) / 255
            
        self.men_word = np.zeros(300)  # "们"的声音
        for i, char in enumerate("们"):
            self.men_word[i] = ord(char) / 255
        
        self.group_visual = np.random.rand(200) * 0.6 + 0.4
        
        # 组合"我们"
        self.wo_men_pattern = np.zeros(300)
        self.wo_men_pattern[:150] = self.wo_word[:150]
        self.wo_men_pattern[150:300] = self.men_word[:150]
        
        # 教会系统识别基本模式
        self.brain.regions['B_shapes'].learn_pattern("knife_shape", self.knife_visual[:150])
        self.brain.regions['B_words'].learn_pattern("knife_word", self.knife_word)
        self.brain.regions['B_words'].learn_pattern("wo_word", self.wo_word)
        self.brain.regions['B_words'].learn_pattern("men_word", self.men_word)
        self.brain.regions['H_hand_motor'].learn_action("withdraw_hand", self.withdraw_pattern)
        
        # 记录初始状态
        self._record_activation()
        
    def _send_stimulus(self):
        """发送刺激到模型"""
        stimulus_type = self.stimulus_type.get()
        content = self.stimulus_content.get()
        
        if not content:
            self._log("请输入刺激内容")
            return
            
        # 根据不同类型准备刺激
        if stimulus_type == "visual":
            # 视觉刺激用随机模式，但保持一定一致性
            # 这里可以使用哈希函数将文本转为固定模式
            np.random.seed(hash(content) % 2**32)
            stimulus = np.random.rand(200) * 0.5 + 0.5
        elif stimulus_type == "auditory":
            # 听觉刺激直接用文本
            stimulus = content
        else:  # touch
            # 触觉刺激也用随机模式
            np.random.seed(hash(content) % 2**32)
            stimulus = np.random.rand(150) * 0.7 + 0.3
            
        # 处理刺激
        activated = self.brain.process_stimulus(stimulus_type, stimulus)
        concepts = self.brain.get_active_concepts()
        actions = self.brain.get_output_actions()
        
        # 记录当前激活
        self._record_activation()
        
        # 更新界面
        self._log(f"处理{stimulus_type}刺激: {content}")
        self._log(f"激活区域: {activated}")
        if concepts:
            self._log(f"激活概念: {concepts}")
        if actions:
            self._log(f"输出动作: {actions}")
            
        self._update_all_plots()
        
    def _run_knife_demo(self):
        """运行刀子学习演示"""
        self._reset_system()
        self._log("=== 刀子学习演示 ===")
        
        # 第一步：看到刀子
        self._log("\n步骤1: 首次看到刀子")
        activated = self.brain.process_stimulus("visual", self.knife_visual)
        self._log(f"激活区域: {activated}")
        self._log(f"当前概念: {self.brain.get_active_concepts()}")
        self._log(f"输出动作: {self.brain.get_output_actions()}")
        self._record_activation()
        self._update_all_plots()
        self.root.update()
        time.sleep(1)
        
        # 第二步：摸刀子被割伤
        self._log("\n步骤2: 摸刀子被割伤")
        pain_stimulus = np.random.rand(150) * 0.7 + 0.3
        activated = self.brain.process_stimulus("touch", pain_stimulus)
        
        # E类区域痛感激活
        self.brain.regions['E_pain'].activate(self.pain_pattern)
        self._log(f"激活区域: {activated}")
        
        # F类区域触发缩手反射
        self.brain.regions['F_pain_withdraw'].activate(np.random.rand(80) * 0.7 + 0.3)
        self.brain.regions['F_pain_withdraw'].regulate(
            self.brain.regions['E_pain'], 
            self.brain.regions['H_hand_motor'], 
            "strengthen"
        )
        
        # 手部执行缩手动作
        self.brain.regions['H_hand_motor'].activate(self.withdraw_pattern)
        self._log(f"输出动作: {self.brain.get_output_actions()}")
        self._record_activation()
        self._update_all_plots()
        self.root.update()
        time.sleep(1)
        
        # 第三步：形成记忆
        self._log("\n步骤3: 形成记忆关联")
        # 形成D区记忆：刀子和痛的关联
        self.brain.regions['D_knife_concept'].activate(self.knife_visual[:150])
        self.brain.regions['D_knife_concept'].associate(self.brain.regions['E_pain'], 0.7)
        self.brain.regions['D_knife_concept'].associate(self.brain.regions['B_shapes'], 0.8)
        
        # 形成G区：看到刀子要避免
        self.brain.regions['G_knife_avoid'].activate(np.random.rand(100) * 0.6 + 0.4)
        self.brain.regions['G_knife_avoid'].learn_regulation(
            self.brain.regions['D_knife_concept'],
            self.brain.regions['H_hand_motor'],
            "strengthen",
            0.5
        )
        self._record_activation()
        self._update_all_plots()
        self.root.update()
        time.sleep(1)
        
        # 第四步：再次看到刀子
        self._log("\n步骤4: 再次看到刀子")
        self.brain.update_all_regions()  # 更新所有区域状态
        activated = self.brain.process_stimulus("visual", self.knife_visual)
        self._log(f"激活区域: {activated}")
        self._log(f"当前概念: {self.brain.get_active_concepts()}")
        self._log(f"输出动作: {self.brain.get_output_actions()}")
        self._record_activation()
        self._update_all_plots()
        
        self._log("\n演示完成: 系统已经学会了看到刀子就回避的反应")
        
    def _run_language_demo(self):
        """运行语言学习演示"""
        self._reset_system()
        self._log("=== 语言学习演示 ===")
        
        # 第一步：分别听到"我"和"们"
        self._log("\n步骤1: 分别听到'我'和'们'")
        activated = self.brain.process_stimulus("auditory", self.wo_word)
        self._log(f"听到'我'激活的区域: {activated}")
        self._log(f"当前概念: {self.brain.get_active_concepts()}")
        self._record_activation()
        self._update_all_plots()
        self.root.update()
        time.sleep(1)
        
        activated = self.brain.process_stimulus("auditory", self.men_word)
        self._log(f"听到'们'激活的区域: {activated}")
        self._log(f"当前概念: {self.brain.get_active_concepts()}")
        self._record_activation()
        self._update_all_plots()
        self.root.update()
        time.sleep(1)
        
        # 创建"我们"的C区域
        if 'C_wo_men' not in self.brain.regions:
            self.brain.regions['C_wo_men'] = brain_model.TypeCRegion("我们组合词", size=300)
        
        # 学习过程
        self._log("\n步骤2: 学习'我们'与群体的关联")
        for i in range(5):
            self._log(f"\n学习循环 {i+1}:")
            
            # 同时看到群体和听到"我们"
            self._log("  看到群体...")
            self.brain.process_stimulus("visual", self.group_visual)
            self._log("  听到'我们'...")
            self.brain.process_stimulus("auditory", self.wo_men_pattern)
            
            # 将"我"和"们"的B区域连接到"我们"的C区域
            self.brain.regions['C_wo_men'].connect_from(self.brain.regions['B_words'])
            
            # 创建复合模式
            self.brain.regions['C_wo_men'].create_compound_pattern([
                self.brain.regions['B_words'] 
            ])
            
            # 将"我们"的C区域与群体的视觉D区域建立联系
            if 'D_group_concept' not in self.brain.regions:
                self.brain.regions['D_group_concept'] = brain_model.TypeDRegion("群体概念", "social", size=150)
            
            self.brain.regions['D_group_concept'].activate(self.group_visual[:150])
            if 'C_wo_men' in self.brain.regions and self.brain.regions['C_wo_men'].is_active():
                self.brain.regions['D_group_concept'].associate(self.brain.regions['C_wo_men'], 0.2 * (i+1))
                self._log(f"  '我们'与群体概念的联系强度: {0.2 * (i+1):.2f}")
            
            # 输出当前活跃的概念
            active_concepts = self.brain.get_active_concepts()
            if active_concepts:
                self._log(f"  当前激活的概念: {active_concepts}")
            else:
                self._log("  当前没有激活的概念")
                
            self._record_activation()
            self._update_all_plots()
            self.root.update()
            
            # 逐渐增强连接
            self.brain.update_all_regions()
            time.sleep(1)
        
        # 最后测试
        self._log("\n步骤3: 测试学习效果")
        # 重置系统状态
        for region in self.brain.regions.values():
            region.activation = np.zeros(region.size)
        
        # 听到"我们"
        activated = self.brain.process_stimulus("auditory", self.wo_men_pattern)
        self._log(f"激活的区域: {activated}")
        self._log(f"当前概念: {self.brain.get_active_concepts()}")
        
        # 验证只听到"我们"是否能激活群体概念
        if self.brain.regions['D_group_concept'].is_active():
            self._log("成功：仅听到'我们'就能激活群体概念！")
        else:
            self._log("学习不完全：仅听到'我们'还不能激活群体概念。")
            
        self._record_activation()
        self._update_all_plots()
        
        self._log("\n演示完成: 系统已经学会了'我们'的概念")
        
    def _record_activation(self):
        """记录当前激活状态"""
        self.current_step += 1
        self.time_steps.append(self.current_step)
        
        for name, region in self.brain.regions.items():
            activation_level = np.max(region.activation) if region.is_active() else 0
            self.activation_history[name].append(activation_level)
            
    def _update_heatmap(self):
        """更新热力图"""
        self.ax_heatmap.clear()
        
        # 获取当前激活值
        region_names = list(self.brain.regions.keys())
        region_types = []
        activation_values = []
        
        for name in region_names:
            region_type = name.split('_')[0]
            region_types.append(region_type)
            activation_values.append(np.max(self.brain.regions[name].activation))
        
        # 创建热力图数据矩阵
        type_set = sorted(set(region_types))
        heatmap_data = np.zeros((len(type_set), len(region_names)))
        
        for i, name in enumerate(region_names):
            region_type = name.split('_')[0]
            type_idx = type_set.index(region_type)
            heatmap_data[type_idx, i] = activation_values[i]
        
        # 使用自定义颜色映射，从蓝色(0)到红色(1)
        cmap = LinearSegmentedColormap.from_list('brain_activation', 
                                                ['blue', 'lightblue', 'yellow', 'red'])
        
        # 绘制热力图
        im = self.ax_heatmap.imshow(heatmap_data, cmap=cmap, aspect='auto', vmin=0, vmax=1)
        
        # 设置坐标轴标签
        self.ax_heatmap.set_xticks(range(len(region_names)))
        self.ax_heatmap.set_xticklabels(region_names, rotation=90, fontsize=8)
        self.ax_heatmap.set_yticks(range(len(type_set)))
        self.ax_heatmap.set_yticklabels(type_set)
        
        # 设置标题
        self.ax_heatmap.set_title("神经区域激活水平")
        
        # 添加颜色条
        plt.colorbar(im, ax=self.ax_heatmap, label="激活水平")
        
        # 更新画布
        self.fig_heatmap.tight_layout()
        self.canvas_heatmap.draw()
        
    def _update_connections(self):
        """更新连接图"""
        self.ax_connections.clear()
        
        # 获取所有区域和连接
        regions = list(self.brain.regions.values())
        
        # 创建节点位置 (简化为圆形布局)
        n_regions = len(regions)
        angles = np.linspace(0, 2*np.pi, n_regions, endpoint=False)
        radius = 10
        pos = {region.name: (radius * np.cos(angle), radius * np.sin(angle)) 
               for region, angle in zip(regions, angles)}
        
        # 绘制所有区域节点
        for region in regions:
            x, y = pos[region.name]
            # 根据激活水平确定颜色
            activation = np.max(region.activation)
            color = plt.cm.plasma(activation)
            
            # 根据激活水平确定大小
            size = 300 * (0.5 + activation) if activation > 0 else 150
            
            # 绘制节点
            self.ax_connections.scatter(x, y, s=size, c=[color], alpha=0.8, edgecolors='black')
            
            # 添加标签
            self.ax_connections.annotate(region.name, (x, y), 
                                       xytext=(0, 0), textcoords='offset points',
                                       ha='center', va='center', fontsize=8)
        
        # 绘制连接
        for source in regions:
            src_x, src_y = pos[source.name]
            
            for target, weight in source.connections.items():
                if hasattr(target, 'name'):  # 确保目标是有效区域
                    tgt_x, tgt_y = pos[target.name]
                    
                    # 根据权重确定线宽
                    lw = 1 + 3 * weight
                    
                    # 根据连接强度确定颜色深浅
                    color = (0, 0, 0, min(1.0, weight * 2))
                    
                    # 绘制箭头
                    self.ax_connections.annotate('', xy=(tgt_x, tgt_y), xytext=(src_x, src_y),
                                              arrowprops=dict(arrowstyle='->', lw=lw, color=color))
        
        # 设置图像属性
        self.ax_connections.set_title("神经区域连接")
        self.ax_connections.set_xlim(-12, 12)
        self.ax_connections.set_ylim(-12, 12)
        self.ax_connections.set_aspect('equal')
        self.ax_connections.axis('off')
        
        # 更新画布
        self.fig_connections.tight_layout()
        self.canvas_connections.draw()
        
    def _update_timeseries(self):
        """更新时间序列图"""
        self.ax_timeseries.clear()
        
        # 选择绘制哪些区域 (A、D、H类区域代表输入、概念和输出)
        selected_regions = []
        for name in self.activation_history:
            if name.startswith('A_') or name.startswith('D_') or name.startswith('H_'):
                selected_regions.append(name)
                
        # 绘制时间序列
        for name in selected_regions:
            if len(self.time_steps) == len(self.activation_history[name]):
                self.ax_timeseries.plot(self.time_steps, self.activation_history[name], 
                                      label=name, marker='o', markersize=3)
        
        # 设置图像属性
        self.ax_timeseries.set_title("神经区域活动时间序列")
        self.ax_timeseries.set_xlabel("时间步")
        self.ax_timeseries.set_ylabel("激活水平")
        self.ax_timeseries.set_ylim(0, 1.1)
        self.ax_timeseries.grid(True, linestyle='--', alpha=0.7)
        self.ax_timeseries.legend(loc='upper left', fontsize=8)
        
        # 更新画布
        self.fig_timeseries.tight_layout()
        self.canvas_timeseries.draw()
        
    def _update_all_plots(self):
        """更新所有绘图"""
        self._update_heatmap()
        self._update_connections()
        self._update_timeseries()
        
    def _reset_system(self):
        """重置系统"""
        # 重新创建脑模型
        self.brain = brain_model.ArtificialBrainModel()
        
        # 重置激活历史
        self.activation_history = {name: [] for name in self.brain.regions}
        self.time_steps = []
        self.current_step = 0
        
        # 重新初始化演示数据
        self._init_demo_data()
        
        # 更新所有绘图
        self._update_all_plots()
        
        # 记录
        self._log("系统已重置")
        
    def _log(self, message):
        """向状态文本框添加日志信息"""
        self.status_text.config(state=tk.NORMAL)
        self.status_text.insert(tk.END, message + "\n")
        self.status_text.see(tk.END)
        self.status_text.config(state=tk.DISABLED)
        
    def _clear_log(self):
        """清空状态文本框"""
        self.status_text.config(state=tk.NORMAL)
        self.status_text.delete(1.0, tk.END)
        self.status_text.config(state=tk.DISABLED)
        

def main():
    """主函数"""
    root = tk.Tk()
    app = BrainVisualizer(root)
    root.mainloop()
    
if __name__ == "__main__":
    main() 